#include "types.h"

#include <cstdio>
#include <unordered_set>
#include <vector>
#include <cmath>
#include <cstring>
#include <cassert>

using namespace std;


void gen_sbox_inside(byte* sbox, 
	int cur_pos, 
	int max_val,
	vector<vector<byte> >& save_des)
{	
	if(cur_pos == max_val)
	{
		save_des.push_back(vector<byte>(sbox, sbox+cur_pos));
		return;
	}
	for(int i = 0 ; i < max_val; ++i)
	{
		sbox[cur_pos] = i;
		gen_sbox_inside(sbox, cur_pos + 1, max_val, save_des);
	}
}

void gen_all_sbox(vector<vector<byte> > &des, int sbox_bits)
{
	int sbox_max_value = (int)pow(2, sbox_bits);
	byte *tmp_sbox = new byte[sbox_max_value];
	
	gen_sbox_inside(tmp_sbox, 0, sbox_max_value, des);
	
	delete [] tmp_sbox;
}

void gen_sbox_possible_diff(vector<vector<byte> > &sbox_all,
	vector<vector<byte> > &possible_diff,
	byte sbox_bits)
{
	int max_val = (int) pow(2, sbox_bits);
	int size = sbox_all.size();
	possible_diff.resize(size);

	for(int i = 0; i < size; ++i)
	{
		vector<byte>& sbox = sbox_all[i];
		vector<byte> &diff_set = possible_diff[i];
		for(int try_in_diff = 0x00; try_in_diff < max_val; ++try_in_diff)
		{
			int found = 0;
			for(int try_out_diff = 0x00; try_out_diff < max_val; ++try_out_diff)
			{
				for(int try_in = 0x00; try_in < max_val; ++try_in)
				{	
					int try_in_e = try_in ^ try_in_diff;
					int sbox_out, sbox_out_e;
					sbox_out = sbox[try_in];
					sbox_out_e = sbox[try_in_e];
					int out_diff = sbox_out ^ sbox_out_e;
					if(out_diff == try_out_diff)
					{
						found = 1;
						diff_set.push_back( (try_in_diff << sbox_bits) | try_out_diff); 
						break;
					}
				}
				if(found)
					break;
			}
		}
	}
}

int main()
{
	const int masks[] = {
		0x01,0x03,0x07,0x0f,0x1f,0x3f,0x7f,0xff
	};
	const int guess_bit_count = 2;
	const int each_cand_count = (int)pow(2, guess_bit_count);
	const int sbox_cand_count = (int)pow(each_cand_count, 
			each_cand_count);
	const int sbox_max_val = each_cand_count;
	const int mask = masks[guess_bit_count - 1];
	vector<vector<byte> > sbox_all;
	vector<vector<byte> > sbox_in_out_diff;
	vector<int> candicates;

	sbox_all.reserve(sbox_cand_count);
	sbox_in_out_diff.resize(sbox_cand_count);

	candicates.reserve(sbox_cand_count);
	for(int i = 0; i < sbox_cand_count; ++i)
	{
		candicates.push_back(i);
	}
	
	gen_all_sbox(sbox_all, guess_bit_count);
	gen_sbox_possible_diff(sbox_all, sbox_in_out_diff, guess_bit_count);

	int count = 0;
	int output_diff, cipher, cipher_e;
	int tmp;//scanf of %x must be 4 byte size
	while(true)
	{
		int ret;
		ret = scanf("%x", &tmp);
		if(ret == EOF)
			break;
		output_diff = tmp;
		ret = scanf("%x", &tmp);
		cipher = tmp;
		ret = scanf("%x", &tmp);
		cipher_e = tmp;

		printf("Try times:%d\n", ++count);

		int try_in, try_adder;
		int found = 0;
	
		unordered_set<int> add_diff_set;
		int try_in_diff;
		for(try_in_diff = 0; try_in_diff < sbox_max_val; ++try_in_diff)
		{
			found = 0;
			for( try_in = 0; try_in < sbox_max_val; ++try_in)
			{
				byte try_in_e = try_in ^ try_in_diff;
				for(try_adder = 0; try_adder < sbox_max_val; ++try_adder)
				{
					byte output,output_e;
					output = (try_in + try_adder) & mask;
					output_e = (try_in_e + try_adder) & mask;
					byte try_output_diff = output ^ output_e;
					if(try_output_diff == output_diff)
					{
						found = 1;
						break;
					}
				}
				if(found)
					break;
			}
			if(found)
				add_diff_set.insert(try_in_diff);
		}
	
		vector<int> tmp_set;

		for(int i = 0; i < candicates.size(); ++i )
		{
			int sbox_pos = candicates[i];
			vector<byte>& sbox_diff_set = sbox_in_out_diff[sbox_pos];
			byte rt,rt_e;

			rt = cipher;
			rt_e = cipher_e;

			byte sbox_in_diff = rt ^ rt_e;

			for(int j = 0; j < sbox_diff_set.size(); ++j)
			{
				int value = sbox_diff_set[j];
				int in_diff = value >> guess_bit_count;
				int out_diff = value & mask;
				if(in_diff == sbox_in_diff)
				{
					if(add_diff_set.find(out_diff) != add_diff_set.end())
					{
						tmp_set.push_back(sbox_pos);
						break;
					}
				}
			}

		}

		candicates = tmp_set;
		printf("Candicates size:%d\n", candicates.size());
		if(candicates.size() == 1)
		{
			printf("Found sbox:\n");

			int found_pos = *candicates.begin();
			vector<byte>& found_sbox = sbox_all[found_pos];
			for(int i = 0; i < found_sbox.size(); ++i)
			{
				printf("%.2x  ", found_sbox[i]);
			}
			printf("\n");
			break;
		}
		else if(candicates.size() == 0)
		{
			printf("Not found any possible sbox!\n");
			return 1;
		}
	}
	return 0;
}

